use super::types::*;
use super::util::*;

use num_traits::identities::{One, Zero};
use std::cmp::min;

/***********
 * DECODER *
 ***********/

#[derive(Debug)]
pub enum DecodeStepResult {
    /// Decoding must go on
    Step,
    /// Finished decoding the current input word
    Finished { choices: Vec<(N, u32)> },
    /// Finished decoding part of the current input word (relevant for prefixfree decoder)
    PrefixParsed {
        choices: Vec<(N, u32)>,
        suffix_index: u32,
    },
}

/// A stack entry, or element, of the decoder state.
#[derive(Clone, Debug)]
pub struct DecoderStateEntry {
    index: u32,
    choices: Vec<(N, u32)>,
    sentential_form: Vec<Symbol>,
}

impl Default for DecoderStateEntry {
    fn default() -> Self {
        Self::new()
    }
}

impl DecoderStateEntry {
    pub fn new() -> DecoderStateEntry {
        DecoderStateEntry {
            index: 0,
            choices: vec![],
            sentential_form: vec![Symbol::Nonterminal(START_SYMBOL)],
        }
    }
    pub fn from_previous_choices(choices: Vec<(N, u32)>) -> DecoderStateEntry {
        DecoderStateEntry {
            index: 0,
            choices,
            sentential_form: vec![Symbol::Nonterminal(START_SYMBOL)],
        }
    }
}

pub fn compute_bounds_lp(grammar: &SCFG, choices: &[(N, u32)]) -> BoundsLP {
    let mut l: Prob = Zero::zero();
    let mut p: Prob = One::one();
    for (n, i) in choices.iter() {
        let rhs = grammar.rules.get(n).unwrap().get(*i as usize).unwrap();
        let l_new = l.clone() + p.clone() * rhs.cumul_lower.clone();
        let p_new = p.clone() * rhs.probability.clone();
        l = l_new;
        p = p_new;
    }

    BoundsLP { l, p }
}

/// Decode via parsing, one step: update the Earley-like data structure
pub fn step_decode(
    grammar: &SCFG,
    configuration: &mut Vec<DecoderStateEntry>,
    input: &[u8],
) -> Result<DecodeStepResult, anyhow::Error> {
    if configuration.is_empty() {
        return Err(anyhow::Error::msg("no parse"));
    }

    let top_entry = configuration.pop().unwrap();

    let replacements = match top_entry.sentential_form.as_slice() {
        [] => {
            if top_entry.index == input.len() as u32 {
                return Ok(DecodeStepResult::Finished {
                    choices: top_entry.choices,
                });
            } else {
                match configuration.len() {
                    0 => {
                        return Err(anyhow::Error::msg(
                            "input prefix parsed, input not exhausted",
                        ))
                    }
                    // There may still be other possibilities
                    _ => {
                        return Ok(DecodeStepResult::Step);
                    }
                }
            }
        }
        [Symbol::Terminals(w), ..] => {
            if input[top_entry.index as usize..].len() < w.len() {
                vec![]
            } else {
                let (w1, _) = input[top_entry.index as usize..].split_at(w.len());
                // Munch or fail
                if w1 == &w[..] {
                    vec![DecoderStateEntry {
                        index: top_entry.index + w.len() as u32,
                        choices: top_entry.choices,
                        sentential_form: top_entry.sentential_form[1..].to_vec(),
                    }]
                } else {
                    vec![]
                }
            }
        }
        &[Symbol::Nonterminal(n), ..] => {
            if let Some(rule_list) = grammar.rules.get(&n) {
                rule_list
                    .iter()
                    .enumerate()
                    .map(|(rule_no, RuleEntry { rhs, .. })| {
                        let mut newx = rhs.clone();
                        newx.append(&mut top_entry.sentential_form[1..].to_vec());

                        let mut choices = top_entry.choices.to_vec();
                        choices.push((n, rule_no as u32));

                        DecoderStateEntry {
                            index: top_entry.index,
                            choices,
                            sentential_form: newx,
                        }
                    })
                    .collect()
            } else {
                return Err(anyhow::Error::msg(format!(
                    "missing rules for {} in grammar",
                    n
                )));
            }
        }
    };

    let mut replacements = replacements;
    replacements.retain(|rep| rep.sentential_form.len() <= input.len() - rep.index as usize);
    configuration.append(&mut replacements);

    Ok(DecodeStepResult::Step)
}

pub fn step_decode_ondemand<'a>(
    grammar: &SCFG,
    configuration: &mut Vec<DecoderStateEntry>,
    collected_input: &mut Vec<u8>,
    remaining_input: &mut (dyn Iterator<Item = u8> + 'a),
) -> Result<DecodeStepResult, anyhow::Error> {
    if configuration.is_empty() {
        return Err(anyhow::Error::msg("no parse"));
    }

    let top_entry = configuration.pop().unwrap();

    let replacements = match top_entry.sentential_form.as_slice() {
        [] => {
            match remaining_input.next() {
                Some(c) => {
                    collected_input.push(c);
                }
                None => {
                    if collected_input.len() <= top_entry.index as usize {
                        return Ok(DecodeStepResult::Finished {
                            choices: top_entry.choices,
                        });
                    }
                }
            };
            return Ok(DecodeStepResult::PrefixParsed {
                choices: top_entry.choices,
                suffix_index: top_entry.index,
            });
        }

        [Symbol::Terminals(w), ..] => {
            let offset = top_entry.index as usize;

            let orig_len = collected_input.len() - offset;
            let smaller_len = min(orig_len, w.len());

            // Fill buffer from input iterator if needed, and compare:
            let matched = w[..smaller_len] == collected_input[offset..offset + smaller_len] && {
                let mut i = orig_len;
                loop {
                    if i >= w.len() {
                        break true;
                    }
                    match remaining_input.next() {
                        Some(c) => {
                            collected_input.push(c);
                            if w[i] != collected_input[offset + i] {
                                break false;
                            }
                        }
                        None => {
                            break false;
                        }
                    };
                    i += 1;
                }
            };

            if matched {
                vec![DecoderStateEntry {
                    index: top_entry.index + w.len() as u32,
                    choices: top_entry.choices,
                    sentential_form: top_entry.sentential_form[1..].to_vec(),
                }]
            } else {
                vec![]
            }
        }
        &[Symbol::Nonterminal(n), ..] => {
            if let Some(rule_list) = grammar.rules.get(&n) {
                rule_list
                    .iter()
                    .enumerate()
                    .rev()
                    .map(|(rule_no, RuleEntry { rhs, .. })| {
                        let mut newx = rhs.clone();
                        newx.append(&mut top_entry.sentential_form[1..].to_vec());

                        let mut choices = top_entry.choices.to_vec();
                        choices.push((n, rule_no as u32));

                        DecoderStateEntry {
                            index: top_entry.index,
                            choices,
                            sentential_form: newx,
                        }
                    })
                    .collect()
            } else {
                return Err(anyhow::Error::msg(format!(
                    "missing rules for {} in grammar",
                    n
                )));
            }
        }
    };

    let mut replacements = replacements;

    replacements.retain(|rep| {
        let len = rep.sentential_form.len() as u32;
        while collected_input.len() as u32 + rep.index < len {
            match remaining_input.next() {
                Some(c) => {
                    collected_input.push(c);
                }
                None => {
                    return false;
                }
            }
        }
        true
    });

    configuration.append(&mut replacements);

    Ok(DecodeStepResult::Step)
}

/// The ordinary Decoder takes a *sequence* of sentences of G as input.
/// The reason is that lower and upper bound may be too far apart after
/// outputting a single sentence, and the encoder must then follow up
/// with another.
pub struct Decoder<'a> {
    grammar: &'a SCFG,
    state: Vec<DecoderStateEntry>,
    input: Vec<u8>,
    remaining_inputs: Box<dyn Iterator<Item = Vec<u8>> + 'a>,
}

/// Here the grammar is prefix-free and the "sentences" are self-delimiting.
pub struct PrefixfreeDecoder<'a> {
    grammar: &'a SCFG,
    state: Vec<DecoderStateEntry>,
    /// The extent of input that has been read from the input_iterator so far
    collected_input: Vec<u8>,
    input_iterator: Box<dyn Iterator<Item = u8> + 'a>,
    /// Did the input cause an error (parse error)?
    /// Then we ought to be able to detect it and handle it.
    error: Option<anyhow::Error>,
}

/// The Prefixfree Decoder as an iterator works more or less like Decoder + LengthDelimitedInputIterator
impl<'a> Iterator for PrefixfreeDecoder<'a> {
    type Item = BoundsLP;

    fn next(&mut self) -> Option<Self::Item> {
        if self.error.is_some() {
            return None;
        }

        loop {
            let k = step_decode_ondemand(
                self.grammar,
                &mut self.state,
                &mut self.collected_input,
                &mut self.input_iterator,
            );

            match k {
                Ok(DecodeStepResult::Step) => {
                    continue;
                }

                Ok(DecodeStepResult::Finished { choices }) => {
                    return Some(compute_bounds_lp(self.grammar, &choices));
                }

                Ok(DecodeStepResult::PrefixParsed {
                    choices,
                    suffix_index,
                }) => {
                    // Unlike Decoder, reset the output every time:
                    // do not assemble and convert it immediately,
                    // otherwise it would not be possible to "trial decode".
                    self.state = vec![DecoderStateEntry {
                        index: 0,
                        choices: vec![],
                        sentential_form: vec![Symbol::Nonterminal(START_SYMBOL)],
                    }];
                    self.collected_input = self.collected_input[suffix_index as usize..].to_vec();

                    return Some(compute_bounds_lp(self.grammar, &choices));
                }

                Err(e) => {
                    self.error = Some(e);
                    return None;
                }
            }
        }
    }
}

/// The Prefixfree Decoder reads from a single stream of u8's.
impl<'a> PrefixfreeDecoder<'a> {
    pub fn new(scfg: &'a SCFG, input: &'a mut dyn Iterator<Item = u8>) -> PrefixfreeDecoder<'a> {
        PrefixfreeDecoder {
            grammar: scfg,
            state: vec![DecoderStateEntry::new()],
            collected_input: vec![],
            input_iterator: Box::new(input),
            error: None,
        }
    }

    pub fn decode(&mut self) -> Result<Vec<u8>, anyhow::Error> {
        let mut everything = vec![];

        loop {
            let k = step_decode_ondemand(
                self.grammar,
                &mut self.state,
                &mut self.collected_input,
                &mut self.input_iterator,
            );

            match k {
                Ok(DecodeStepResult::Step) => {
                    continue;
                }

                Ok(DecodeStepResult::Finished { choices }) => {
                    let l = compute_bounds_lp(self.grammar, &choices).l;
                    let mut part = binaryfraction_between_0_and_1_to_bytes(&l);
                    everything.append(&mut part);
                    return Ok(everything);
                }

                Ok(DecodeStepResult::PrefixParsed {
                    choices,
                    suffix_index,
                }) => {
                    self.state = vec![DecoderStateEntry {
                        index: 0,
                        choices,
                        sentential_form: vec![Symbol::Nonterminal(START_SYMBOL)],
                    }];
                    self.collected_input = self.collected_input[suffix_index as usize..].to_vec();
                }

                Err(e) => {
                    return Err(e);
                }
            }
        }
    }
}

/// The plain decoder, without the requirement for the grammar to be prefix-free.
/// Instead, it relies on the sentences being delimited externally.
impl<'a> Decoder<'a> {
    pub fn new(scfg: &'a SCFG, inputs: &'a mut dyn Iterator<Item = Vec<u8>>) -> Decoder<'a> {
        let head = inputs.next().unwrap();

        Decoder {
            grammar: scfg,
            state: vec![DecoderStateEntry::new()],
            input: head,
            remaining_inputs: Box::new(inputs),
        }
    }

    /// @precondition: use only when remaining_input is empty
    /* pub fn refill(&'a mut self, additional_inputs: &'a mut dyn Iterator<Item = Vec<u8>>) {
        let head = additional_inputs.next().unwrap();
        self.input = head;
        self.remaining_inputs = Box::new(additional_inputs);
    } */

    pub fn decode(&mut self) -> Result<BoundsLP, anyhow::Error> {
        loop {
            let k = step_decode(self.grammar, &mut self.state, &self.input);

            match k {
                Ok(DecodeStepResult::Step) => {
                    continue;
                }
                Ok(DecodeStepResult::Finished { choices }) => {
                    // in case we want to refill with additional inputs
                    // TODO: test whether this works as intended; apply
                    self.state = vec![DecoderStateEntry::from_previous_choices(choices.clone())];
                    match self.remaining_inputs.next() {
                        None => {
                            let bounds_lp = compute_bounds_lp(self.grammar, &choices);
                            return Ok(bounds_lp);
                        }
                        Some(input) => {
                            let next = input;
                            self.input = next;
                        }
                    }
                }

                // Since the grammar is unambiguous, and this is the parser that must
                // parse the entire input (language not assumed to be prefix-free),
                // it must be fatal.
                Ok(DecodeStepResult::PrefixParsed { .. }) => {
                    return Err(anyhow::Error::msg("Only a prefix was parsed"));
                }

                Err(e) => {
                    return Err(e);
                }
            }
        }
    }
}
